#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/5/25 10:05
@desc: 
"""
import itk
import numpy as np

from utils.config import Number


def re_sample(image, output_spacing,
              output_size=None,
              interpolation='linear',
              fill_value=Number.BACKGROUND_PIXEL):
    output_spacing = np.array(output_spacing)
    input_spacing = np.array(image.GetSpacing())
    region = image.GetLargestPossibleRegion()
    input_size = np.array(region.GetSize())
    if output_size is None:
        output_size = input_spacing * input_size / output_spacing
        output_size = np.round(output_size).astype(np.int64)
    sampler = itk.ResampleImageFilter[itk.Image.SS3, itk.Image.SS3].New()
    transform = itk.IdentityTransform[itk.D, 3].New()
    transform.SetIdentity()
    if interpolation == 'linear':
        func = itk.LinearInterpolateImageFunction
    else:
        func = itk.NearestNeighborInterpolateImageFunction
    interpolator = func[itk.Image.SS3, itk.D].New()
    sampler.SetInterpolator(interpolator)
    sampler.SetInput(image)
    sampler.SetTransform(transform)
    sampler.SetOutputOrigin(image.GetOrigin())
    sampler.SetOutputSpacing(output_spacing.tolist())
    sampler.SetOutputDirection(image.GetDirection())
    sampler.SetSize(output_size.tolist())
    sampler.SetDefaultPixelValue(fill_value)
    sampler.UpdateLargestPossibleRegion()
    return sampler.GetOutput()
